
import os,matplotlib
import numpy as np
import matplotlib.pyplot as plt

def readrain(fname):
    infile = open(fname,'r')
    temp = infile.readline()

    while not 'NATOM' in temp:
        temp = infile.readline()
    temp = infile.readline().split() # vals
    natom, nmol, nreact, nfix, nvary = int(temp[0]), int(temp[1]), int(temp[2]), int(temp[4]), int(temp[6])
    temp = infile.readline() # header      
    temp = infile.readline().split() # vals
    nz = int(temp[2])

    # skip all lines that come before 'MIXING RATIO' 
    temp = infile.readline()
    while 'REACTION RATES' not in temp:
        temp = infile.readline()

    while not '139)' in temp:
        temp = infile.readline()
    while not 'COL' in temp:
        temp = infile.readline()
    xx = temp.split()
    r139 = float(xx[9])
    r140 = float(xx[10])
    temp = infile.readline()
    while not '142)' in temp:
        temp = infile.readline()
    temp = infile.readline()
    while not 'COL' in temp:
        temp = infile.readline()
    xx = temp.split()
    r142 = float(xx[2])
    dat = [r139,r140,r142]
    print(fname,dat)

    temp = infile.readline()
    while not 'FLUX:' in temp:
        temp = infile.readline()
    while not 'HNO2' in temp:
        temp = infile.readline()
    temp = infile.readline().split()
    print(dat)
    dat[0] += -1*float(temp[8])
    dat[1] += -1*float(temp[9])
    dat[2] += -1*float(temp[10])
    print(dat)
#    exit()

    return dat

mypath = ['outputs_0x','outputs_1x','outputs_3x','outputs_10x']
dat = {}
for i in mypath:
    dat[i] = {}
    for ifile in os.listdir('./withsep/'+i+'/'):
        if not ifile.endswith('mini'):
            continue
        dat[i][ifile] = readrain('./withsep/'+i+'/'+ifile)

matplotlib.rc('font', family='serif')
matplotlib.rc('xtick', labelsize = 18)
matplotlib.rc('ytick', labelsize = 18)
matplotlib.rc('axes', linewidth=2)
fontname = 'serif'
fontsize=30
lw = 3.3

pre = ['7mbar','50mbar','200mbar','500mbar','1bar']
fig,ax=plt.subplots(ncols=1,nrows=1,figsize=(15,6))


no2,no3,no4 = np.zeros((4,5)),np.zeros((4,5)),np.zeros((4,5))
for i in mypath:
    if  '_0x' in i:
        ictr = 0
    elif '1x' in i:
        ictr = 1
    elif '3x' in i:
        ictr = 2
    elif '10x' in i:
        ictr = 3
    else:
        print('ice case not found')

    for ifile in list(dat[i].keys()):
        kctr = 0

        if pre[0] in ifile:
            jctr = 0
        elif pre[1] in ifile:
            jctr = 1
        elif pre[2] in ifile:
            jctr = 2
        elif pre[3] in ifile:
            jctr = 3
        elif pre[4] in ifile:
            jctr = 4
        else:
            print('pressure not found, ifile')
        
        no2[ictr,jctr] = dat[i][ifile][0]
        no3[ictr,jctr] = dat[i][ifile][1]
        no4[ictr,jctr] = dat[i][ifile][2]
        print(ictr,jctr,no2[ictr,jctr])
        
lns = ['-','--',':']
for i in range(1,4):
    ax.semilogy(np.arange(5),no2[i,:],color='royalblue',linestyle=lns[i-1],linewidth=2)
    ax.semilogy(np.arange(5),no3[i,:],color='darkorange',linestyle=lns[i-1],linewidth=2)
    ax.semilogy(np.arange(5),no4[i,:],color='mediumvioletred',linestyle=lns[i-1],linewidth=2)

for i in range(0,4):
    ax.scatter(np.arange(5),no2[i,:],color='royalblue',marker='o')#,markersize=10)
    ax.scatter(np.arange(5),no3[i,:],color='darkorange',marker='o')#,markersize=10)
    ax.scatter(np.arange(5),no4[i,:],color='mediumvioletred',marker='o')#,markersize=10)

ax.tick_params('both',length=10,width=2,which='major')
ax.tick_params('both',length=5,width=1,which='minor')
ax.tick_params(axis='x',pad=8)
ax.tick_params(axis='y',pad=8)

plt.savefig('rainrates.png',dpi=300)
